
#include "./driver.h"


namespace lukv {

driver::driver(int _id) : id(_id), sock(socket(AF_INET, SOCK_STREAM, 0)), paddr(&manager::driver_addrs[_id]) { }

driver::~driver() {
    close(sock);
}

#ifdef __OSX__
int driver::start() {
    const int BACKLOG_SIZE = 16;
    const int READ_BUF_SIZE = 2048;
    const int K_EVENTS_CNT = 128;

    set_fd_noblk(sock);

    if (bind(sock, (sockaddr *)paddr, sizeof(*paddr)) < 0) {
        close_fds(sock);
        EXIT_ERR(strerror(errno));
    }

    if (listen(sock, BACKLOG_SIZE)) {
        close_fds(sock);
        EXIT_ERR(strerror(errno));
    }

    int kq = kqueue();

    struct kevent read_event { (uint64_t) sock, EVFILT_READ, EV_ADD | EV_ENABLE, 0, 0, nullptr };

    if (kevent(kq, &read_event, 1, nullptr, 0, nullptr) < 0) {
        close_fds(sock, kq);
        EXIT_ERR(strerror(errno));
    }

    struct kevent out_events[K_EVENTS_CNT];

    char read_buf[READ_BUF_SIZE];

    timespec over_t { 0, 500000000 };

    while (true) {
        int cnt = kevent(kq, nullptr, 0, out_events, K_EVENTS_CNT, &over_t);
        if (cnt < 0) {
            LOG_ERR(strerror(errno));
            break;
        }
        for (int i = 0; i < cnt; ++i) {
            auto _sock = out_events[i].ident;
            auto _data = out_events[i].data;

            if (out_events[i].filter & EVFILT_READ) {
                println("[READ_EV]", "[ident:", _sock, ']', "[data:", _data, ']');

                if (_sock == sock) {
                    sockaddr_in cli_addr { };
                    socklen_t  cli_addr_len = sizeof(cli_addr);

                    for (int j = 0; j < _data; ++j) {
                        int cli_sock = accept(sock, (sockaddr *) &cli_addr, &cli_addr_len);
                        if (cli_sock < 0) {
                            LOG_ERR(strerror(errno));
                            continue;
                        }

                        set_fd_noblk(cli_sock);

                        println("[sock:", cli_sock, ']', "[client:", inet_ntoa(cli_addr.sin_addr), ':', ntohs(cli_addr.sin_port), ']');

                        if (!manager::is_node(cli_addr)) manager::set_expire_cli(cli_sock);

                        struct kevent cli_rw_events { (uintptr_t) cli_sock, EVFILT_READ | EVFILT_WRITE, EV_ENABLE | EV_ADD, 0, 0, nullptr };
                        kevent(kq, &cli_rw_events, 1, nullptr, 0, nullptr);
                    }
                }
                else {
                    auto rlen = recv(_sock, read_buf, _data, 0);

                    if (rlen > 0) {
                        std::string read_str(read_buf, rlen);
                        println("[recv]<", read_str, '>');

                        std::string wt_msg = "D" + handle_cmd(read_str);

                        send(_sock, wt_msg.c_str(), wt_msg.size(), 0);
                        println("[send]<", wt_msg, '>');

                        if (read_str == "#exit") {
                            close_fds(_sock);
                            struct kevent cli_event { (uintptr_t) _sock, EVFILT_VNODE, EV_DISABLE | EV_DELETE, 0, 0, nullptr };
                            kevent(kq, &cli_event, 1, nullptr, 0, nullptr);
                            manager::remove_expire_cli(_sock);
                            println("closed fd: ", _sock);
                            continue;
                        }
                    }
                    else {
                        close_fds(_sock);
                        struct kevent cli_event { (uintptr_t) _sock, EVFILT_VNODE, EV_DISABLE | EV_DELETE, 0, 0, nullptr };
                        kevent(kq, &cli_event, 1, nullptr, 0, nullptr);
                        manager::remove_expire_cli(_sock);
                        println("closed fd: ", _sock);
                        continue;
                    }

                }
            }
        }

        std::vector<int> expired = manager::get_expired_clis();
        for (auto & cli: expired) {
            close_fds(cli);
            manager::remove_expire_cli(cli);
        }

    }

    return 0;
}
#endif

#ifdef __LINUX__
int driver::start() {

    const int EP_EVENT_CNT = 16;
    const int BACK_LOG_SIZE = 32;
    const int READ_BUF_SIZE = 1024;

    set_fd_noblk(sock);

    if (bind(sock, (sockaddr *)paddr, sizeof(*paddr)) < 0) {
        close_fds(sock);
        EXIT_ERR(strerror(errno));
    }

    if (listen(sock, BACK_LOG_SIZE) < 0) {
        close_fds(sock);
        EXIT_ERR(strerror(errno));
    }

    int ep_fd = epoll_create(64);

    struct epoll_event srv_event { };
    srv_event.events = EPOLLIN;
    srv_event.data.fd = sock;

    epoll_ctl(ep_fd, EPOLL_CTL_ADD, sock, &srv_event);

    struct epoll_event ep_events[EP_EVENT_CNT];

    char read_buf[READ_BUF_SIZE];

    while (true) {
        int cnt = epoll_wait(ep_fd, ep_events, EP_EVENT_CNT, 500);
        if (cnt < 0) {
            LOG_ERR(strerror(errno));
            break;
        }
        
        for (int i = 0; i < cnt; i++) {
            int _sock = ep_events[i].data.fd;

            if (_sock == sock) {
                struct sockaddr_in cli_addr { };
                socklen_t cli_addr_len = sizeof(cli_addr);

                int cli_sock = accept(sock, (sockaddr *)&cli_addr, &cli_addr_len);
                if (cli_sock < 0) {
                    LOG_ERR(strerror(errno));
                    continue;
                }

                set_fd_noblk(cli_sock);

                if (!manager::is_node(cli_addr)) manager::set_expire_cli(cli_sock);

                println("[sock:", cli_sock, ']', "[client:", inet_ntoa(cli_addr.sin_addr), ':', ntohs(cli_addr.sin_port), ']');

                srv_event.events = EPOLLIN | EPOLLET;
                srv_event.data.fd = cli_sock;

                epoll_ctl(ep_fd, EPOLL_CTL_ADD, cli_sock, &srv_event);
            }
            else {
                char * _ptr_r = read_buf;

                int r_len = recv(_sock, read_buf, READ_BUF_SIZE, 0);

                if (r_len <= 0) {
                    close_fds(_sock);
                    epoll_ctl(ep_fd, EPOLL_CTL_DEL, _sock, nullptr);
                    manager::remove_expire_cli(_sock);
                    println("closed fd: ", _sock);
                    continue;
                }
                else {
                    _ptr_r += r_len;
                    while (r_len = recv(_sock, _ptr_r, READ_BUF_SIZE, 0) > 0) {
                        _ptr_r += r_len;
                    }
                }

                std::string r_msg(read_buf, _ptr_r);
                println("[recv]<", r_msg, '>');

                std::string w_msg = "D" + handle_cmd(r_msg);
                send(_sock, w_msg.c_str(), w_msg.size(), 0);
                println("[send]<", w_msg, '>');

                if (r_msg == "#exit") {
                    close_fds(_sock);
                    epoll_ctl(ep_fd, EPOLL_CTL_DEL, _sock, nullptr);
                    manager::remove_expire_cli(_sock);
                    println("closed fd: ", _sock);
                    continue;
                }
            }
        }

        std::vector<int> expired = manager::get_expired_clis();
        for (auto & cli: expired) {
            close_fds(cli);
            manager::remove_expire_cli(cli);
        }
    }

    return 0;
}
#endif

std::string driver::handle_cmd(const std::string & _cmd) {
    // 1. 来自客户端：将命令发给worker执行，若执行成功，则通知其它driver相应该事件，修改key的元数据
    if (_cmd[0] == 'C') {
        std::string r_cmd = _cmd.substr(1);
        std::vector<std::string> args = split_by_spaces(r_cmd);
        if (args.size() < 3) return "EERROR";
        std::string & type = args[0];
        std::string & opt = args[1];
        std::string & key = args[2];
        auto worker_cnt = manager::worker_addrs.size();
        std::vector<int> worker_ids = random_ints_unique(0, worker_cnt - 1, manager::replication);
        char read_buf[2048];
        std::string ans;
        for (auto wid : worker_ids) {
            if (connect_to_worker(wid) < 0) continue;
            write(cluster_worker_socks[wid], _cmd.c_str(), _cmd.size());
            if (read_buf[0] == 'S') {
                if (opt != "del") key_workers[key].insert(wid);
                std::string _dtmp;
                _dtmp.append("D").append(key).append(" ").append(std::to_string(wid)).append(" ").append(opt);
                send_to_other_drivers(_dtmp);
            }
            auto r_len = read(cluster_worker_socks[wid], read_buf, 2048);
            ans.append(std::string(read_buf, r_len)).append("|");
        }
        return ans;
    }
        // 2. 来自其它driver（可通过机器id或ip地址判断是否为本机）：根据发过来的通知解析事件，并执行相应操作（修改key元数据、删除key等）
    else if (_cmd[0] == 'D') {
        std::string r_cmd = _cmd.substr(1);
        std::vector<std::string> args = split_by_spaces(r_cmd);
        if (args.size() < 3) return "EERROR";
        std::string & key = args[0];
        int wid = std::stoi(args[1]);
        std::string & opt = args[2];
        if (opt == "del") {
            key_workers[key].erase(wid);
            if (key_workers[key].empty()) {
                key_workers.erase(key);
            }
        }
        else {
            key_workers[key].insert(wid);
        }
        return "SSUCC";
    }
        // 3. 来自worker：直接发送到客户端
    else if (_cmd[0] == 'W') {
        return _cmd.substr(1);
    }
    else {
        return "ERROR";
    }
}

int driver::connect_to_worker(int wid) {
    if (cluster_worker_socks.find(wid) == cluster_worker_socks.end()) {
        sockaddr_in worker_addr = manager::worker_addrs[wid];
        int cur_sock = socket(AF_INET, SOCK_STREAM, 0);
        set_fd_noblk(cur_sock);
        cluster_worker_socks[wid] = cur_sock;
        if (connect(cur_sock, (sockaddr *)&worker_addr, sizeof(worker_addr)) < 0) {
            close_fds(cur_sock);
            LOG_ERR(strerror(errno));
            return -1;
        }
    }
    return 0;
}

int driver::connect_to_driver(int did) {
    if (cluster_driver_socks.find(did) == cluster_driver_socks.end()) {
        sockaddr_in driver_addr = manager::driver_addrs[did];
        int cur_sock = socket(AF_INET, SOCK_STREAM, 0);
        set_fd_noblk(cur_sock);
        cluster_driver_socks[did] = cur_sock;
        if (connect(cur_sock, (sockaddr *)&driver_addr, sizeof(driver_addr)) < 0) {
            close_fds(cur_sock);
            LOG_ERR(strerror(errno));
            return -1;
        }
    }
    return 0;
}

int driver::send_to_other_drivers(const std::string & msg) {
    for (auto i = 0; i < manager::driver_addrs.size(); ++i) {
        if (i == id) continue;
        connect_to_driver(i);
        write(cluster_driver_socks[i], msg.c_str(), msg.size());
    }
    return 0;
}

}
